package batcher

import "time"

// Batcher 是批量处理的通用结构体
type Batcher[T any] struct {
	size         int           // 批量大小
	multiple     int           // 批量大小倍数，用于调整缓冲区大小
	interval     time.Duration // 定时写入间隔
	callback     func([]T)     // 用户提供的回调函数，传递批量数据
	dataChannel  chan T        // 数据通道
	closeChannel chan struct{} // 用于关闭的通道
}

// NewBatcher 初始化一个新的 Batcher
func NewBatcher[T any](size, multiple int, interval time.Duration, callback func([]T)) *Batcher[T] {
	if size < 1 {
		size = 1
	}
	if multiple < 1 {
		multiple = 1
	}
	if interval < time.Millisecond {
		interval = time.Millisecond
	}
	b := &Batcher[T]{
		size:         size,
		multiple:     multiple,
		interval:     interval,
		callback:     callback,
		dataChannel:  make(chan T, size*multiple),
		closeChannel: make(chan struct{}),
	}
	go b.worker()
	return b
}

// Push 向批量处理队列中添加数据
func (b *Batcher[T]) Push(data T) {
	b.dataChannel <- data
}

// Close 关闭 Batcher 并在关闭前处理所有剩余数据
func (b *Batcher[T]) Close() {
	close(b.closeChannel)
}

// worker 启动批量处理逻辑
func (b *Batcher[T]) worker() {
	var buffer []T
	ticker := time.NewTicker(b.interval)
	defer ticker.Stop()

	for {
		select {
		case data := <-b.dataChannel:
			buffer = append(buffer, data)
			if len(buffer) >= b.size {
				b.callback(buffer)
				buffer = buffer[:0]
				ticker.Reset(b.interval)
			}

		case <-ticker.C:
			if len(buffer) > 0 {
				b.callback(buffer)
				buffer = buffer[:0]
			}

		case <-b.closeChannel:
			if len(buffer) > 0 {
				b.callback(buffer)
			}
			return
		}
	}
}
